{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ardha27/AI-Song-Cover-RVC/blob/main/TrainingV2_NoUI.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "[![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/R6R7AH1FA)\n",
        "<a href=\"https://trakteer.id/ardha27\">\n",
        "    <img src=\"https://cdn.trakteer.id/images/embed/trbtn-red-1.png\" alt=\"Trakteer\" height=\"35\">\n",
        "</a>"
      ],
      "metadata": {
        "id": "Dy2HJED2ya_i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Mount Drive\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "jwu07JgqoFON",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Unzip File From Drive\n",
        "\n",
        "!unzip -n /content/drive/MyDrive/assets/KiwKiw.zip -d /content\n",
        "%cd /content/KiwKiw\n",
        "!pip install pip==23.3.1\n",
        "!pip install -r requirements.txt -q"
      ],
      "metadata": {
        "id": "ge_97mfpgqTm",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "# @title Create Model Folder\n",
        "model_name = \"kafka\" # @param {type:\"string\"}\n",
        "exp_dir = model_name\n",
        "dataset = \"/content/drive/MyDrive/dataset/kafka\" # @param {type:\"string\"}\n",
        "sample_rate = \"48000\" # @param [\"40000\", \"48000\"]\n",
        "ksample_rate = \"48k\"\n",
        "if sample_rate == \"40000\":\n",
        "  ksample_rate = \"40k\"\n",
        "else:\n",
        "  ksample_rate = \"48k\"\n",
        "version = \"v2\" # @param [\"v1\", \"v2\"]\n",
        "version19 = version\n",
        "\n",
        "f0method = \"rmvpe_gpu\" # @param [\"pm\", \"dio\", \"harvest\", \"rmvpe\", \"rmvpe_gpu\"]\n",
        "\n",
        "save_frequency = 50 # @param {type:\"slider\", min:0, max:50, step:1}\n",
        "epoch = 10 # @param {type:\"integer\"}\n",
        "batch_size = \"7\" # @param {type:\"string\"}\n",
        "cache_gpu = True # @param {type:\"boolean\"}\n",
        "\n",
        "now_dir = \"/content/KiwKiw\"\n",
        "\n",
        "os.makedirs(\"%s/logs/%s\" % (now_dir, exp_dir), exist_ok=True)\n",
        "f = open(\"%s/logs/%s/preprocess.log\" % (now_dir, exp_dir), \"w\")\n",
        "os.makedirs(\"%s/logs/%s\" % (now_dir, exp_dir), exist_ok=True)\n",
        "f = open(\"%s/logs/%s/extract_f0_feature.log\" % (now_dir, exp_dir), \"w\")\n",
        "f.close()"
      ],
      "metadata": {
        "id": "QIJLtXk2KRn5",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Process Data\n",
        "command = f\"python infer/modules/train/preprocess.py '{dataset}' {sample_rate} 2 '{now_dir}/logs/{exp_dir}' False 3.0\"\n",
        "print(command)\n",
        "!{command}"
      ],
      "metadata": {
        "id": "EAY0Hz2pI_bW",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Feature Extraction\n",
        "\n",
        "if f0method != \"rmvpe_gpu\":\n",
        "  command = f\"python infer/modules/train/extract/extract_f0_print.py '{now_dir}/logs/{exp_dir}' 2 '{f0method}'\"\n",
        "else:\n",
        "  command = f\"python infer/modules/train/extract/extract_f0_rmvpe.py 1 0 0 '{now_dir}/logs/{exp_dir}' True\"\n",
        "print(command)\n",
        "!{command}\n",
        "\n",
        "command = f\"python infer/modules/train/extract_feature_print.py cuda:0 1 0 0 '{now_dir}/logs/{exp_dir}' '{version}' False\"\n",
        "print(command)\n",
        "!{command}"
      ],
      "metadata": {
        "id": "Q1qr2RBnDe7k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Train Feature Index\n",
        "import numpy as np\n",
        "import faiss\n",
        "\n",
        "\n",
        "def train_index(exp_dir1, version19):\n",
        "    exp_dir = \"logs/%s\" % (exp_dir1)\n",
        "    os.makedirs(exp_dir, exist_ok=True)\n",
        "    feature_dir = (\n",
        "        \"%s/3_feature256\" % (exp_dir)\n",
        "        if version19 == \"v1\"\n",
        "        else \"%s/3_feature768\" % (exp_dir)\n",
        "    )\n",
        "    if not os.path.exists(feature_dir):\n",
        "        return \"请先进行特征提取!\"\n",
        "    listdir_res = list(os.listdir(feature_dir))\n",
        "    if len(listdir_res) == 0:\n",
        "        return \"请先进行特征提取！\"\n",
        "    infos = []\n",
        "    npys = []\n",
        "    for name in sorted(listdir_res):\n",
        "        phone = np.load(\"%s/%s\" % (feature_dir, name))\n",
        "        npys.append(phone)\n",
        "    big_npy = np.concatenate(npys, 0)\n",
        "    big_npy_idx = np.arange(big_npy.shape[0])\n",
        "    np.random.shuffle(big_npy_idx)\n",
        "    big_npy = big_npy[big_npy_idx]\n",
        "    if big_npy.shape[0] > 2e5:\n",
        "        infos.append(\"Trying doing kmeans %s shape to 10k centers.\" % big_npy.shape[0])\n",
        "        yield \"\\n\".join(infos)\n",
        "        try:\n",
        "            big_npy = (\n",
        "                MiniBatchKMeans(\n",
        "                    n_clusters=10000,\n",
        "                    verbose=True,\n",
        "                    batch_size=256 * config.n_cpu,\n",
        "                    compute_labels=False,\n",
        "                    init=\"random\",\n",
        "                )\n",
        "                .fit(big_npy)\n",
        "                .cluster_centers_\n",
        "            )\n",
        "        except:\n",
        "            info = traceback.format_exc()\n",
        "            logger.info(info)\n",
        "            infos.append(info)\n",
        "            yield \"\\n\".join(infos)\n",
        "\n",
        "    np.save(\"%s/total_fea.npy\" % exp_dir, big_npy)\n",
        "    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)\n",
        "    infos.append(\"%s,%s\" % (big_npy.shape, n_ivf))\n",
        "    yield \"\\n\".join(infos)\n",
        "    index = faiss.index_factory(256 if version19 == \"v1\" else 768, \"IVF%s,Flat\" % n_ivf)\n",
        "    # index = faiss.index_factory(256if version19==\"v1\"else 768, \"IVF%s,PQ128x4fs,RFlat\"%n_ivf)\n",
        "    infos.append(\"training\")\n",
        "    yield \"\\n\".join(infos)\n",
        "    index_ivf = faiss.extract_index_ivf(index)  #\n",
        "    index_ivf.nprobe = 1\n",
        "    index.train(big_npy)\n",
        "    faiss.write_index(\n",
        "        index,\n",
        "        \"%s/trained_IVF%s_Flat_nprobe_%s_%s_%s.index\"\n",
        "        % (exp_dir, n_ivf, index_ivf.nprobe, exp_dir1, version19),\n",
        "    )\n",
        "\n",
        "    infos.append(\"adding\")\n",
        "    yield \"\\n\".join(infos)\n",
        "    batch_size_add = 8192\n",
        "    for i in range(0, big_npy.shape[0], batch_size_add):\n",
        "        index.add(big_npy[i : i + batch_size_add])\n",
        "    faiss.write_index(\n",
        "        index,\n",
        "        \"%s/added_IVF%s_Flat_nprobe_%s_%s_%s.index\"\n",
        "        % (exp_dir, n_ivf, index_ivf.nprobe, exp_dir1, version19),\n",
        "    )\n",
        "    infos.append(\n",
        "        \"成功构建索引，added_IVF%s_Flat_nprobe_%s_%s_%s.index\"\n",
        "        % (n_ivf, index_ivf.nprobe, exp_dir1, version19)\n",
        "    )\n",
        "\n",
        "result_generator = train_index(exp_dir, version)\n",
        "\n",
        "for result in result_generator:\n",
        "    print(result)"
      ],
      "metadata": {
        "id": "xzurT_PYQOIJ",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Train Model\n",
        "\n",
        "from random import shuffle\n",
        "import json\n",
        "import os\n",
        "import pathlib\n",
        "from subprocess import Popen, PIPE, STDOUT\n",
        "\n",
        "# Remove the logging setup\n",
        "\n",
        "def click_train(\n",
        "    exp_dir1,\n",
        "    sr2,\n",
        "    if_f0_3,\n",
        "    spk_id5,\n",
        "    save_epoch10,\n",
        "    total_epoch11,\n",
        "    batch_size12,\n",
        "    if_save_latest13,\n",
        "    pretrained_G14,\n",
        "    pretrained_D15,\n",
        "    gpus16,\n",
        "    if_cache_gpu17,\n",
        "    if_save_every_weights18,\n",
        "    version19,\n",
        "):\n",
        "    # 生成filelist\n",
        "    exp_dir = \"%s/logs/%s\" % (now_dir, exp_dir1)\n",
        "    os.makedirs(exp_dir, exist_ok=True)\n",
        "    gt_wavs_dir = \"%s/0_gt_wavs\" % (exp_dir)\n",
        "    feature_dir = (\n",
        "        \"%s/3_feature256\" % (exp_dir)\n",
        "        if version19 == \"v1\"\n",
        "        else \"%s/3_feature768\" % (exp_dir)\n",
        "    )\n",
        "    if if_f0_3:\n",
        "        f0_dir = \"%s/2a_f0\" % (exp_dir)\n",
        "        f0nsf_dir = \"%s/2b-f0nsf\" % (exp_dir)\n",
        "        names = (\n",
        "            set([name.split(\".\")[0] for name in os.listdir(gt_wavs_dir)])\n",
        "            & set([name.split(\".\")[0] for name in os.listdir(feature_dir)])\n",
        "            & set([name.split(\".\")[0] for name in os.listdir(f0_dir)])\n",
        "            & set([name.split(\".\")[0] for name in os.listdir(f0nsf_dir)])\n",
        "        )\n",
        "    else:\n",
        "        names = set([name.split(\".\")[0] for name in os.listdir(gt_wavs_dir)]) & set(\n",
        "            [name.split(\".\")[0] for name in os.listdir(feature_dir)]\n",
        "        )\n",
        "    opt = []\n",
        "    for name in names:\n",
        "        if if_f0_3:\n",
        "            opt.append(\n",
        "                \"%s/%s.wav|%s/%s.npy|%s/%s.wav.npy|%s/%s.wav.npy|%s\"\n",
        "                % (\n",
        "                    gt_wavs_dir.replace(\"\\\\\", \"\\\\\\\\\"),\n",
        "                    name,\n",
        "                    feature_dir.replace(\"\\\\\", \"\\\\\\\\\"),\n",
        "                    name,\n",
        "                    f0_dir.replace(\"\\\\\", \"\\\\\\\\\"),\n",
        "                    name,\n",
        "                    f0nsf_dir.replace(\"\\\\\", \"\\\\\\\\\"),\n",
        "                    name,\n",
        "                    spk_id5,\n",
        "                )\n",
        "            )\n",
        "        else:\n",
        "            opt.append(\n",
        "                \"%s/%s.wav|%s/%s.npy|%s\"\n",
        "                % (\n",
        "                    gt_wavs_dir.replace(\"\\\\\", \"\\\\\\\\\"),\n",
        "                    name,\n",
        "                    feature_dir.replace(\"\\\\\", \"\\\\\\\\\"),\n",
        "                    name,\n",
        "                    spk_id5,\n",
        "                )\n",
        "            )\n",
        "    fea_dim = 256 if version19 == \"v1\" else 768\n",
        "    if if_f0_3:\n",
        "        for _ in range(2):\n",
        "            opt.append(\n",
        "                \"%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s/logs/mute/2a_f0/mute.wav.npy|%s/logs/mute/2b-f0nsf/mute.wav.npy|%s\"\n",
        "                % (now_dir, sr2, now_dir, fea_dim, now_dir, now_dir, spk_id5)\n",
        "            )\n",
        "    else:\n",
        "        for _ in range(2):\n",
        "            opt.append(\n",
        "                \"%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s\"\n",
        "                % (now_dir, sr2, now_dir, fea_dim, spk_id5)\n",
        "            )\n",
        "    shuffle(opt)\n",
        "    with open(\"%s/filelist.txt\" % exp_dir, \"w\") as f:\n",
        "        f.write(\"\\n\".join(opt))\n",
        "\n",
        "    # Replace logger.debug, logger.info with print statements\n",
        "    print(\"Write filelist done\")\n",
        "    print(\"Use gpus:\", str(gpus16))\n",
        "    if pretrained_G14 == \"\":\n",
        "        print(\"No pretrained Generator\")\n",
        "    if pretrained_D15 == \"\":\n",
        "        print(\"No pretrained Discriminator\")\n",
        "    if version19 == \"v1\" or sr2 == \"40k\":\n",
        "        config_path = \"configs/v1/%s.json\" % sr2\n",
        "    else:\n",
        "        config_path = \"configs/v2/%s.json\" % sr2\n",
        "    config_save_path = os.path.join(exp_dir, \"config.json\")\n",
        "    if not pathlib.Path(config_save_path).exists():\n",
        "        with open(config_save_path, \"w\", encoding=\"utf-8\") as f:\n",
        "            with open(config_path, \"r\") as config_file:\n",
        "                config_data = json.load(config_file)\n",
        "                json.dump(\n",
        "                    config_data,\n",
        "                    f,\n",
        "                    ensure_ascii=False,\n",
        "                    indent=4,\n",
        "                    sort_keys=True,\n",
        "                )\n",
        "            f.write(\"\\n\")\n",
        "\n",
        "    cmd = (\n",
        "        'python infer/modules/train/train.py -e \"%s\" -sr %s -f0 %s -bs %s -g %s -te %s -se %s %s %s -l %s -c %s -sw %s -v %s'\n",
        "        % (\n",
        "            exp_dir1,\n",
        "            sr2,\n",
        "            1 if if_f0_3 else 0,\n",
        "            batch_size12,\n",
        "            gpus16,\n",
        "            total_epoch11,\n",
        "            save_epoch10,\n",
        "            \"-pg %s\" % pretrained_G14 if pretrained_G14 != \"\" else \"\",\n",
        "            \"-pd %s\" % pretrained_D15 if pretrained_D15 != \"\" else \"\",\n",
        "            1 if if_save_latest13 == True else 0,\n",
        "            1 if if_cache_gpu17 == True else 0,\n",
        "            1 if if_save_every_weights18 == True else 0,\n",
        "            version19,\n",
        "        )\n",
        "    )\n",
        "    # Use PIPE to capture the output and error streams\n",
        "    p = Popen(cmd, shell=True, cwd=now_dir, stdout=PIPE, stderr=STDOUT, bufsize=1, universal_newlines=True)\n",
        "\n",
        "    # Print the command's output as it runs\n",
        "    for line in p.stdout:\n",
        "        print(line.strip())\n",
        "\n",
        "    # Wait for the process to finish\n",
        "    p.wait()\n",
        "    return \"训练结束, 您可查看控制台训练日志或实验文件夹下的train.log\"\n",
        "\n",
        "if version == 'v1':\n",
        "    if ksample_rate == '40k':\n",
        "        G_path = 'assets/pretrained/f0G40k.pth'\n",
        "        D_path = 'assets/pretrained/f0D40k.pth'\n",
        "    elif ksample_rate == '48k':\n",
        "        G_path = 'assets/pretrained/f0G48k.pth'\n",
        "        D_path = 'assets/pretrained/f0D48k.pth'\n",
        "elif version == 'v2':\n",
        "    if ksample_rate == '40k':\n",
        "        G_path = 'assets/pretrained_v2/f0G40k.pth'\n",
        "        D_path = 'assets/pretrained_v2/f0D40k.pth'\n",
        "    elif ksample_rate == '48k':\n",
        "        G_path = 'assets/pretrained_v2/f0G48k.pth'\n",
        "        D_path = 'assets/pretrained_v2/f0D48k.pth'\n",
        "\n",
        "result_generator = click_train(\n",
        "    exp_dir,\n",
        "    ksample_rate,\n",
        "    True,\n",
        "    0,\n",
        "    save_frequency,\n",
        "    epoch,\n",
        "    batch_size,\n",
        "    True,\n",
        "    G_path,\n",
        "    D_path,\n",
        "    0,\n",
        "    cache_gpu,\n",
        "    False,\n",
        "    version,\n",
        ")\n",
        "print(result_generator)\n",
        "\n"
      ],
      "metadata": {
        "id": "rIreV7qsa3Es",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Copy Training Result to Drive\n",
        "!mkdir -p /content/drive/MyDrive/Model/{exp_dir}\n",
        "\n",
        "!cp /content/KiwKiw/logs/{exp_dir}/added_*.index /content/drive/MyDrive/Model/{exp_dir}\n",
        "!cp /content/KiwKiw/logs/{exp_dir}/total_*.npy /content/drive/MyDrive/Model/{exp_dir}\n",
        "!cp /content/KiwKiw/assets/weights/{exp_dir}.pth /content/drive/MyDrive/Model/{exp_dir}\n",
        "\n",
        "%cd /content/drive/MyDrive/Model/{exp_dir}\n",
        "!zip {exp_dir}.zip added_*.index total_*.npy {exp_dir}.pth\n",
        "%cd /content/KiwKiw"
      ],
      "metadata": {
        "id": "FgJuNeAwx5Y_",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}